#include "parser.h"

// 删除文件读取函数
// 使用传入的 tokens 列表

void Parser::error(const string& msg) {
    if (pos < tokens.size()) {
        cerr << "Syntax error: " << msg << " at token: " << tokens[pos].lexeme
            << " (line " << tokens[pos].line << ", column " << tokens[pos].column << ")" << endl;
    }
    else {
        cerr << "Syntax error: " << msg << " at end of file" << endl;
    }
    exit(1);
}

void Parser::match(TokenType expected) {
    if (pos < tokens.size() && tokens[pos].type == expected) {
        pos++;
    }
    else {
        string expectedStr;
        switch (expected) {
        case TokenType::INT: expectedStr = "INT"; break;
        case TokenType::VOID: expectedStr = "VOID"; break;
        case TokenType::IF: expectedStr = "IF"; break;
        case TokenType::ELSE: expectedStr = "ELSE"; break;
        case TokenType::WHILE: expectedStr = "WHILE"; break;
        case TokenType::BREAK: expectedStr = "BREAK"; break;
        case TokenType::CONTINUE: expectedStr = "CONTINUE"; break;
        case TokenType::RETURN: expectedStr = "RETURN"; break;
        case TokenType::IDENTIFIER: expectedStr = "IDENTIFIER"; break;
        case TokenType::NUMBER: expectedStr = "NUMBER"; break;
        case TokenType::ASSIGN: expectedStr = "ASSIGN"; break;
        case TokenType::OR: expectedStr = "OR"; break;
        case TokenType::AND: expectedStr = "AND"; break;
        case TokenType::LT: expectedStr = "LT"; break;
        case TokenType::GT: expectedStr = "GT"; break;
        case TokenType::LE: expectedStr = "LE"; break;
        case TokenType::GE: expectedStr = "GE"; break;
        case TokenType::EQ: expectedStr = "EQ"; break;
        case TokenType::NE: expectedStr = "NE"; break;
        case TokenType::ADD: expectedStr = "ADD"; break;
        case TokenType::SUB: expectedStr = "SUB"; break;
        case TokenType::MUL: expectedStr = "MUL"; break;
        case TokenType::DIV: expectedStr = "DIV"; break;
        case TokenType::MOD: expectedStr = "MOD"; break;
        case TokenType::NOT: expectedStr = "NOT"; break;
        case TokenType::COMMA: expectedStr = "COMMA"; break;
        case TokenType::SEMICOLON: expectedStr = "SEMICOLON"; break;
        case TokenType::LPAREN: expectedStr = "LPAREN"; break;
        case TokenType::RPAREN: expectedStr = "RPAREN"; break;
        case TokenType::LBRACE: expectedStr = "LBRACE"; break;
        case TokenType::RBRACE: expectedStr = "RBRACE"; break;
        case TokenType::END_OF_FILE: expectedStr = "END_OF_FILE"; break;
        default: expectedStr = "UNKNOWN"; break;
        }

        string actualStr = (pos < tokens.size()) ?
            to_string(static_cast<int>(tokens[pos].type)) : "END_OF_FILE";

        error("Expected " + expectedStr + ", got " + actualStr);
    }
}

unique_ptr<ASTNode> Parser::PrimaryExpr() {
    auto node = make_unique<ExprNode>("Primary");
    
    // 处理标识符（可能是变量或函数调用）
    if (pos < tokens.size() && tokens[pos].type == TokenType::IDENTIFIER) {
        string identName = tokens[pos].lexeme;
        node->children.emplace_back(make_unique<ExprNode>(identName));
        match(TokenType::IDENTIFIER);
        
        // 如果是函数调用必须跟括号
        if (pos < tokens.size() && tokens[pos].type == TokenType::LPAREN) {
            match(TokenType::LPAREN);
            
            // 处理参数列表
            if (pos < tokens.size() && tokens[pos].type != TokenType::RPAREN) {
                node->children.push_back(Expr());
                while (pos < tokens.size() && tokens[pos].type == TokenType::COMMA) {
                    match(TokenType::COMMA);
                    node->children.push_back(Expr());
                }
            }
            
            // 确保有右括号
            if (pos >= tokens.size() || tokens[pos].type != TokenType::RPAREN) {
                error("Expected ')' after function arguments");
            } else {
                match(TokenType::RPAREN);
            }
        }
        // 如果不是函数调用，后面不能直接跟数字
        else if (pos < tokens.size() && tokens[pos].type == TokenType::NUMBER) {
            error("Function call requires parentheses");
        }
    }
    // 处理数字字面量
    else if (pos < tokens.size() && tokens[pos].type == TokenType::NUMBER) {
        node->children.emplace_back(make_unique<NumberNode>(tokens[pos].lexeme, false));
        match(TokenType::NUMBER);
    }
    // 处理括号表达式
    else if (pos < tokens.size() && tokens[pos].type == TokenType::LPAREN) {
        match(TokenType::LPAREN);
        node->children.push_back(Expr());
        match(TokenType::RPAREN);
    }
    else {
        error("Expected identifier, number or parenthesized expression");
    }
    
    return node;
}

unique_ptr<ASTNode> Parser::AddExpr() {
    auto node = MulExpr();
    while (pos < tokens.size() &&
           (tokens[pos].type == TokenType::ADD ||
            tokens[pos].type == TokenType::SUB)) {
        
        auto opNode = make_unique<ExprNode>(tokens[pos].lexeme);
        match(tokens[pos].type);
        opNode->children.push_back(move(node));
        opNode->children.push_back(MulExpr());
        node = move(opNode);
    }
    return node;
}

unique_ptr<ASTNode> Parser::UnaryExpr() {
    if (pos < tokens.size() &&
        (tokens[pos].type == TokenType::ADD ||
         tokens[pos].type == TokenType::SUB ||
         tokens[pos].type == TokenType::NOT)) {
        
        auto node = make_unique<ExprNode>(tokens[pos].lexeme);
        match(tokens[pos].type);
        node->children.push_back(UnaryExpr());
        return node;
    } else {
        return PrimaryExpr();
    }
}

unique_ptr<ASTNode> Parser::MulExpr() {
    auto node = UnaryExpr();
    while (pos < tokens.size() &&
           (tokens[pos].type == TokenType::MUL ||
            tokens[pos].type == TokenType::DIV ||
            tokens[pos].type == TokenType::MOD)) {
        
        auto opNode = make_unique<ExprNode>(tokens[pos].lexeme);
        match(tokens[pos].type);
        opNode->children.push_back(move(node));
        opNode->children.push_back(UnaryExpr());
        node = move(opNode);
    }
    return node;
}

unique_ptr<ASTNode> Parser::RelExpr() {
    auto node = AddExpr();
    while (pos < tokens.size() &&
           (tokens[pos].type == TokenType::LT ||
            tokens[pos].type == TokenType::GT ||
            tokens[pos].type == TokenType::LE ||
            tokens[pos].type == TokenType::GE ||
            tokens[pos].type == TokenType::EQ ||
            tokens[pos].type == TokenType::NE)) {
        
        auto opNode = make_unique<ExprNode>(tokens[pos].lexeme);
        match(tokens[pos].type);
        opNode->children.push_back(move(node));
        opNode->children.push_back(AddExpr());
        node = move(opNode);
    }
    return node;
}

unique_ptr<ASTNode> Parser::LAndExpr() {
    auto node = RelExpr();
    while (pos < tokens.size() && tokens[pos].type == TokenType::AND) {
        auto opNode = make_unique<ExprNode>(tokens[pos].lexeme);
        match(TokenType::AND);
        opNode->children.push_back(move(node));
        opNode->children.push_back(RelExpr());
        node = move(opNode);
    }
    return node;
}

unique_ptr<ASTNode> Parser::LOrExpr() {
    auto node = LAndExpr();
    while (pos < tokens.size() && tokens[pos].type == TokenType::OR) {
        auto opNode = make_unique<ExprNode>(tokens[pos].lexeme);
        match(TokenType::OR);
        opNode->children.push_back(move(node));
        opNode->children.push_back(LAndExpr());
        node = move(opNode);
    }
    return node;
}


unique_ptr<ASTNode> Parser::Expr() {
    return LOrExpr();
}

unique_ptr<ASTNode> Parser::Param() {
    auto node = make_unique<ExprNode>("Param");
    match(TokenType::INT);
    node->children.emplace_back(make_unique<ExprNode>(tokens[pos].lexeme));
    match(TokenType::IDENTIFIER);
    return node;
}

unique_ptr<ASTNode> Parser::FuncDef() {
    string retType;
    if (tokens[pos].type == TokenType::INT) {
        retType = "int";
    } else if (tokens[pos].type == TokenType::VOID) {
        retType = "void";
    } else {
        error("Expected return type (int or void)");
    }
    match(tokens[pos].type);
    
    string funcName = tokens[pos].lexeme;
    match(TokenType::IDENTIFIER);
    
    auto funcNode = make_unique<FuncDefNode>(retType, funcName);
    match(TokenType::LPAREN);
    
    // 处理参数列表
    if (pos < tokens.size() && tokens[pos].type == TokenType::INT) {
        funcNode->params.push_back(Param());
        while (pos < tokens.size() && tokens[pos].type == TokenType::COMMA) {
            match(TokenType::COMMA);
            funcNode->params.push_back(Param());
        }
    }
    
    match(TokenType::RPAREN);
    funcNode->body = Block();
    return funcNode;
}

unique_ptr<ASTNode> Parser::Block() {
    auto blockNode = make_unique<StmtNode>("Block");
    match(TokenType::LBRACE);
    while (pos < tokens.size() && tokens[pos].type != TokenType::RBRACE) {
        blockNode->children.push_back(Stmt());
    }
    match(TokenType::RBRACE);
    return blockNode;
}

unique_ptr<ASTNode> Parser::Stmt3() {
    unique_ptr<StmtNode>stmtNode;
    stmtNode = make_unique<StmtNode>("ExprStmt");
    stmtNode->children.push_back(Expr());
    match(TokenType::SEMICOLON);
    return stmtNode;
}

unique_ptr<ASTNode> Parser::Stmt4() {
    unique_ptr<StmtNode>stmtNode;
    stmtNode = make_unique<StmtNode>("AssignStmt");
    stmtNode->children.emplace_back(make_unique<ExprNode>(tokens[pos].lexeme));
    match(TokenType::IDENTIFIER);
    match(TokenType::ASSIGN);
    stmtNode->children.push_back(Expr());
    match(TokenType::SEMICOLON);
    return stmtNode;
}

unique_ptr<ASTNode> Parser::Stmt5() {
    unique_ptr<StmtNode>stmtNode;
    stmtNode = make_unique<StmtNode>("DeclareStmt");

    match(TokenType::INT);
    stmtNode->children.emplace_back(make_unique<ExprNode>(tokens[pos].lexeme));
    match(TokenType::IDENTIFIER);
    if (pos < tokens.size() && tokens[pos].type == TokenType::ASSIGN) {
        match(TokenType::ASSIGN);
        stmtNode->children.push_back(Expr());
    }
    match(TokenType::SEMICOLON);
    return stmtNode;
}

unique_ptr<ASTNode> Parser::Stmt6() {
    unique_ptr<StmtNode> stmtNode;
    stmtNode = make_unique<StmtNode>("IfStmt");
    match(TokenType::IF);
    
    // 确保有左括号
    if (tokens[pos].type != TokenType::LPAREN) {
        error("Expected '(' after 'if'");
    }
    match(TokenType::LPAREN);
    
    stmtNode->children.push_back(Expr());
    
    // 确保有右括号
    if (tokens[pos].type != TokenType::RPAREN) {
        error("Expected ')' after if condition");
    }
    match(TokenType::RPAREN);
    
    stmtNode->children.push_back(Stmt());
    
    if (pos < tokens.size() && tokens[pos].type == TokenType::ELSE) {
        match(TokenType::ELSE);
        stmtNode->children.push_back(Stmt());
    }
    
    return stmtNode;
}

unique_ptr<ASTNode> Parser::Stmt7() {
    unique_ptr<StmtNode> stmtNode;
    stmtNode = make_unique<StmtNode>("WhileStmt");
    match(TokenType::WHILE);
    
    // 确保有左括号
    if (tokens[pos].type != TokenType::LPAREN) {
        error("Expected '(' after 'while'");
    }
    match(TokenType::LPAREN);
    
    stmtNode->children.push_back(Expr());
    
    // 确保有右括号
    if (tokens[pos].type != TokenType::RPAREN) {
        error("Expected ')' after while condition");
    }
    match(TokenType::RPAREN);
    
    stmtNode->children.push_back(Stmt());
    return stmtNode;
}

unique_ptr<ASTNode> Parser::Stmt10() {
    unique_ptr<StmtNode>stmtNode;
    stmtNode = make_unique<StmtNode>("ReturnStmt");
    match(TokenType::RETURN);
    if (pos < tokens.size() && tokens[pos].type != TokenType::SEMICOLON) {
        stmtNode->children.push_back(Expr());
    }
    match(TokenType::SEMICOLON);
    return stmtNode;
}

unique_ptr<ASTNode> Parser::Stmt() {
    unique_ptr<ASTNode>stmtNode;

    if (pos >= tokens.size()) {
        error("Unexpected end of file in Stmt");
    }

    if (tokens[pos].type == TokenType::LBRACE) {
        stmtNode = Block();
    }
    else if (tokens[pos].type == TokenType::SEMICOLON) {
        stmtNode = make_unique<StmtNode>("EmptyStmt");
        match(TokenType::SEMICOLON);
    }
    else if ((tokens[pos].type == TokenType::IDENTIFIER &&
        pos + 1 < tokens.size() &&
        tokens[pos + 1].type != TokenType::ASSIGN) ||
        tokens[pos].type == TokenType::NUMBER ||
        tokens[pos].type == TokenType::LPAREN ||
        (tokens[pos].type == TokenType::IDENTIFIER &&
            pos + 1 < tokens.size() &&
            tokens[pos + 1].type == TokenType::LPAREN) ||
        tokens[pos].type == TokenType::ADD ||
        tokens[pos].type == TokenType::SUB ||
        tokens[pos].type == TokenType::NOT) {
        stmtNode = Stmt3();
    }
    else if (tokens[pos].type == TokenType::IDENTIFIER &&
        pos + 1 < tokens.size() &&
        tokens[pos + 1].type == TokenType::ASSIGN) {
        stmtNode = Stmt4();
    }
    else if (tokens[pos].type == TokenType::INT) {
        stmtNode = Stmt5();
    }
    else if (tokens[pos].type == TokenType::IF) {
        stmtNode = Stmt6();
    }
    else if (tokens[pos].type == TokenType::WHILE) {
        stmtNode = Stmt7();
    }
    else if (tokens[pos].type == TokenType::BREAK) {
        stmtNode = make_unique<StmtNode>("BreakStmt");
        match(TokenType::BREAK);
        match(TokenType::SEMICOLON);
    }
    else if (tokens[pos].type == TokenType::CONTINUE) {
        stmtNode = make_unique<StmtNode>("ContinueStmt");
        match(TokenType::CONTINUE);
        match(TokenType::SEMICOLON);
    }
    else if (tokens[pos].type == TokenType::RETURN) {
        stmtNode = Stmt10();
    }
    else {
        error("Unexpected token in Stmt");
    }
    return stmtNode;
}

unique_ptr<ASTNode> Parser::CompUnit() {
    auto compNode = make_unique<CompUnitNode>();
    compNode->funcs.push_back(FuncDef());
    while (pos < tokens.size() && tokens[pos].type != TokenType::END_OF_FILE) {
        compNode->funcs.push_back(FuncDef());
    }
    return compNode;
}